package cas.ihep.spark.io;

import java.io.*;
import java.nio.ByteBuffer;
import java.util.NoSuchElementException;

public class Transformer {

    public interface Source {
        ByteBuffer next() throws Exception;
        long index();
        long count();
    }

    public interface Sink extends Closeable{
        void put(ByteBuffer buf)throws Exception;
    }

    public interface MultiSink extends Sink{
        Sink currentSink()throws Exception;
        default void put(ByteBuffer buf)throws Exception{
            currentSink().put(buf);
        }
    }

    public static Sink fromOutputStream(OutputStream os,boolean needCache){
        return new OutputStreamSink(os,needCache);
    }

    public static Source fromInputStream(InputStream is,int eleSize,int eleCount){
        return new InputStreamSource(is,eleSize,eleCount);
    }

    public static Source fromByteBuffer(ByteBuffer buf,int eleSize,int eleCount){
        if(eleCount*eleSize<=buf.remaining()){
            return new ByteBufferSource(buf,eleSize,eleCount);
        }
        throw new IllegalArgumentException();
    }

    @SuppressWarnings("WeakerAccess")
    public static class OutputStreamSink implements Sink, Closeable {
        OutputStream output;
        public OutputStreamSink(OutputStream os,boolean needBuffered){
            if(needBuffered && ! (os instanceof BufferedOutputStream)){
                output=new BufferedOutputStream(os);
            }else{
                output=os;
            }
        }
        @Override
        public void put(ByteBuffer buf) throws Exception {
            int rem=buf.remaining();
            if(buf.hasArray()){
                byte[] hb=buf.array();
                int pos=buf.position();
                output.write(hb,pos,rem);
            }else{
                byte[] tmp=new byte[rem];
                buf.get(tmp);
                output.write(tmp);
            }
        }

        @Override
        public void close() throws IOException {
            output.close();
        }
    }

    private static class InputStreamSource implements Source,Closeable{

        InputStream input;
        byte[] hb;
        int index;
        int count;

        InputStreamSource(InputStream is,int s,int c){
            if(is instanceof BufferedInputStream){
                input=is;
            }else{
                input=new BufferedInputStream(is);
            }
            hb=new byte[s];
            index=0;
            count=c;
        }

        @Override
        public ByteBuffer next() throws Exception {
            int ret=input.read(hb);
            if(ret<hb.length){
                throw new NoSuchElementException();
            }
            index++;
            return ByteBuffer.wrap(hb);
        }

        @Override
        public long index() {
            return index;
        }

        @Override
        public long count() {
            return count;
        }

        @Override
        public void close() throws IOException {
            input.close();
        }
    }

    private static class ByteBufferSource implements Source{
        ByteBuffer buffer;
        int index;
        int count;
        int size;
        ByteBufferSource(ByteBuffer buf,int eleSize,int eleCount){
            buffer=buf;
            index=0;
            count=eleCount;
            size=eleSize;
        }

        public ByteBuffer next(){
            if(index<count){
                return ByteBufferUtil.sliceReadOnly(buffer,(int)(index++)*size,size);
            }
            throw new NoSuchElementException();
        }
        public long index(){
            return index;
        }
        public long count(){
            return count;
        }
    }

    public static void transform(Source input, Sink output)throws Exception{
        while(input.index()<input.count()){
            output.put(input.next());
        }
    }
}
